import numpy
import math
from dolfin import *
from common import directional_diameter

def coth(x):
    return 1./math.tanh(x)       

def get_tau(v,eps,p_degree=1):
    '''Compute stabilization parameter for SUPG. Parameter is defined following Galeao et all: FE analysis of conv.-diff.-react...'''
    mesh = v.function_space().mesh()
    limit = mesh.hmax()
    nc = mesh.num_cells()
    taus = numpy.zeros(nc)

    # solve problem vMag*phi = sqrt(inner(v,v)) to get the magnitude of velocity
    V = FunctionSpace(mesh,'CG',v.element().degree())
    vNorm = TrialFunction(V)
    vNormTest = TestFunction(V)

    a = vNorm*vNormTest*dx 
    l = sqrt(inner(v,v))*vNormTest*dx

    vNorm = Function(V)
    solve(a==l,vNorm)

    for i in xrange(nc):
        c = Cell(mesh,i)
        
        x = c.midpoint()
        v_mag = vNorm(x.x(),x.y()) # interpolate to barycenter to get cell's velocity magnitude 

        if v_mag < 1E-16:
            taus[i] = 0 
        else:
            h = directional_diameter(c,v)
            
            eps = float(eps)                              # to avoid some ufl.* magic
            Pe = v_mag*h/2./eps/p_degree                  # cell's Peclet number
            tau = h/2./v_mag/p_degree*(coth(Pe)-1./Pe)    # optimal stab. parameter: coth(Pe)-1./Pe; critical: max(0,1.-1./Pe) see. do Carmo(2004)
            taus[i] = tau
            
            if (h-limit)>0: # this should never get raised. with h=directional_diameter(), h is bounded by hypotenuse   
                raise ValueError('Value of element diameter is wrong!') 
                
    return taus

